import torch
from torchmetrics import Metric
from myloss import generate_mask


class InstanceAccuracy(Metric):

    def __init__(self, dist_sync_on_step=False):
        super().__init__(dist_sync_on_step=dist_sync_on_step)

        self.add_state("correct", default=torch.tensor(0),
                       dist_reduce_fx="sum")
        self.add_state("total", default=torch.tensor(0), dist_reduce_fx="sum")

    def update(self, logits: torch.Tensor, labels: torch.Tensor, lengths: list[int]):
        """更新状态

        Args:
            logits: 模型预测的logits, [total_ocr_len].
            labels: dataloader加载的一维tensor, [total_ocr_len, ].
            lengths: 指代每个ocr对应的长度.
        """
        logits = (logits > 0).int()
        count = 0
        for l in lengths:
            if torch.all(logits[count: count+l] == labels[count: count+l]):
                self.correct += 1
            count += l
        self.total += len(lengths)
    
    def compute(self):
        return self.correct.float() / self.total


class RankingAccuracy(Metric):
    """统计匹配的准确率(平均以及完美匹配)
    """
    def __init__(self, dist_sync_on_step=False):
        super().__init__(dist_sync_on_step=dist_sync_on_step)

        self.add_state("correct", default=torch.tensor(0),
                       dist_reduce_fx="sum")
        self.add_state("correct_pair", default=torch.tensor(0),
                       dist_reduce_fx="sum")
        self.add_state("total_pair", default=torch.tensor(0), dist_reduce_fx="sum")
        self.add_state("total", default=torch.tensor(0), dist_reduce_fx="sum")
    

    def update(self, logits: torch.Tensor, rankorder: torch.Tensor, lengths: list[int]):
        """更新状态

        Args:
            logits: 模型预测的logits, [batch, seq, 2].
            rankorder: dataloader加载的一维tensor, [total_ocr_len, ].
            lengths: 指代每个ocr对应的长度.
        """
        rankscore = logits[:, :, 1]
        mask = (logits[:, :, 0] < 0)
        seq_order = -torch.ones(logits.size()[:2]).type_as(logits)
        # seq_order [batch, seq]
        count = 0
        for i in range(len(lengths)):
            seq_order[i, :lengths[i]] = rankorder[count: count+lengths[i]]
            count += lengths[i]
        seq_order[mask] = -1
        valid_len = torch.sum(seq_order >= 0, dim=1)
        # 按正确顺序进行排列
        seq_order[seq_order < 0] = logits.size(1) + 1
        _, true_order = torch.sort(seq_order)
        rankscore = rankscore.gather(1, true_order)
        # 计算得分差
        rankdiff = torch.diff(rankscore)
        diff_mask = generate_mask(valid_len-1, logits.size(1)-1)
        # 得分高的应该排前面，diff为负数
        correct_count = rankdiff < 0
        correct_count[diff_mask == 0] = 0
        correct_count = torch.sum(correct_count, dim=1)
        perfect_correct_num = torch.sum(correct_count == (valid_len-1))

        self.total += len(lengths)
        self.correct += perfect_correct_num

    
    def compute(self):
        perfect_accuracy = self.correct / self.total
        return perfect_accuracy

            
